# data/hit_aggregator.py

import pandas as pd


def aggregate_by_stock_date(hits_df: pd.DataFrame) -> pd.DataFrame:
    """
    对 3s 级 hits 数据，按 stock_code + add_date 聚合出统计特征。
    返回包含字段：
    - sum_money_gt_300w: money >= 300万的累计成交额
    - money_count_ge_300w: money >= 3%（300w）的出现次数
    - money_count_ge_300w: money >= 5%（500w）的出现次数
    """

    if hits_df.empty:
        return pd.DataFrame(columns=[
            "stock_code", "add_date",
            "sum_money_gt_300w",
            "money_count_ge_300w",
            "money_count_ge_500w"
        ])

    grouped = hits_df.groupby(["stock_code", "add_date"]).agg(
        sum_money_gt_300w=pd.NamedAgg(
            column="money",
            aggfunc=lambda x: x[x >= 300].sum()
        ),
        money_count_ge_300w=pd.NamedAgg(
            column="money",
            aggfunc=lambda x: (x >= 300).sum()
        ),
        money_count_ge_500w=pd.NamedAgg(
            column="money",
            aggfunc=lambda x: (x >= 500).sum()
        )
    ).reset_index()

    # 字段命名统一为 features.py 中定义的字段名
    grouped.rename(columns={
        "money_count_ge_300w": "money_count_ge_300w",
        "money_count_ge_500w": "money_count_ge_500w"
    }, inplace=True)

    return grouped


# def aggregate_by_stock_date(hits_df: pd.DataFrame, money_threshold: float = 0.01) -> pd.DataFrame:
#     """
#     聚合 hits 数据为个股-日期级别特征，用于训练。

#     参数：
#         hits_df: 原始 hits 数据（包含 stock_code、add_date、money 字段）
#         money_threshold: 判定“放量波动”的 money 阈值（单位：亿元）

#     返回：
#         DataFrame，字段包括：
#             - stock_code
#             - add_date
#             - sum_money: 当日总成交额（money字段加总）
#             - fluctuation_count: 大于阈值的成交次数（波动次数）
#             - avg_money: 平均单次成交额
#             - max_money: 当日最大成交额
#     """
#     if hits_df.empty:
#         return pd.DataFrame()

#     grouped = hits_df.groupby(["stock_code", "add_date"])

#     agg_df = grouped.agg(
#         sum_money=("money", "sum"),
#         fluctuation_count=("money", lambda x: (x >= money_threshold).sum()),
#         avg_money=("money", "mean"),
#         max_money=("money", "max")
#     ).reset_index()

#     return agg_df
